from data.gazebo import GazeboTrain, GazeboValidate, GazeboTest
from data.mdcd import MDCDTestMars, Haworth
from models.detect import (
    BandeiraDetector,
    EmamiHighlightShadowDetector,
    EmamiHoughDetector,
)
import cv2
import torch
from torchvision.utils import draw_bounding_boxes
import os

output_dir = "/disk527/sdb1/a804_cbf/datasets/harworth"
if __name__ == "__main__":
    if not os.path.exists(os.path.join(output_dir, "bbox")):
        os.makedirs(os.path.join(output_dir, "bbox"))
    # Load the data
    val_data = Haworth(db_dir="/disk527/sdb1/a804_cbf/datasets/harworth")

    # first stage processed data
    first_stage = EmamiHoughDetector(median_size=3, th=10)

    for i, data in enumerate(val_data):
        if i < 100:
            continue
        image = (
            torch.tensor(data["raw"])
            .unsqueeze(0)
            .float()
            .unsqueeze(0)
            .repeat(1, 3, 1, 1)
        )
        # iamge = image / 127.5 - 1
        image = (image - image.min()) / (image.max() - image.min()) * 2 - 1
        for img, bbox in zip(image, first_stage(image)):
            if bbox.shape[0] == 0:
                continue
            img = ((img + 1) * 127.5).to(dtype=torch.uint8)
            x_bbox = draw_bounding_boxes(img, bbox, colors=(240, 10, 157))
            x_bbox = x_bbox.permute(1, 2, 0).cpu().numpy()
            cv2.imwrite(os.path.join(output_dir, "bbox", f"{i}.png"), x_bbox)
